import numpy as np
# import cPickle
import os
import h5py
from itertools import chain
from easy_io import read_pkl_file, write_pkl_file

thre_score = 0.1
thre_dist = 10
thre_xy = 5


# nms_para: nms or not 
# spac: spacing of vol 
def nms(center, score, thre_dist, spac, nms_para):
    num_candis = len(center)
    if nms_para:

        dists = np.zeros([num_candis, num_candis])
        for r in range(num_candis):
            for c in range(num_candis):
                dists[r][c] = ((center[r][0] - center[c][0]) * spac[0]) ** 2 + \
                              ((center[r][1] - center[c][1]) * spac[1]) ** 2 + ((center[r][2] - center[c][2]) * spac[
                    2]) ** 2

        order = sorted(range(num_candis), key=lambda k: score[k], reverse=True)
        keep = []
        while len(order) > 0:
            i = order[0]
            keep.append(i)
            inds = []
            for j in range(len(order)):
                if dists[i][order[j]] > thre_dist ** 2:
                    inds.append(j)

            tmp = []
            for i_ind in range(len(inds)):
                tmp.append(order[inds[i_ind]])
            order = tmp


    else:
        keep = range(num_candis)

    rf_center = []
    for ir in range(len(keep)):
        rf_center.append(center[keep[ir]])

    return keep, rf_center


def evaluate(center, radius, lab, spac):
    candisPos = []
    candisNeg = []
    isFound = np.zeros(len(lab))
    if len(lab) > 0:
        for icandis in range(len(center)):
            flag = 0
            for ilab in range(len(lab)):
                can = center[icandis]
                cen = lab[ilab]
                if ((can[0] - cen[0]) * spac[0]) ** 2 + ((can[1] - cen[1]) * spac[1]) ** 2 + (
                            (can[2] - cen[2]) * spac[2]) ** 2 < (2 + radius[ilab]) ** 2:
                    isFound[ilab] = 1
                    candisPos.append(can)
                    flag = 1
            if flag == 0:
                candisNeg.append(center[icandis])
    else:
        candisNeg = center
    return len(candisNeg), len(candisPos), isFound


def main_nms(candis):
    scan_id_tmp = candis[0]['scanid']
    source_tmp = candis[0]['source']
    final_candis = []
    center_all = []
    score_all = []
    index_all = []
    iscan = 0

    # num_pos_t=0
    # num_neg_t=0
    # find_nod=0
    # find_ma=0
    # all_nod=0
    # all_ma=0


    # l1 = np.load('/data-disk/Kaggle/backup/lidc/label_dict.npy').item()
    # l2 = np.load('/data-disk/Kaggle/backup/kaggle/label_dict_0325.npy').item()
    # v1 = h5py.File('/data-disk/Kaggle/backup/lidc/vol.hdf5', 'r')
    # v2 = h5py.File('/data-disk/Kaggle/backup/kaggle/vol.hdf5', 'r')
    v = {
        'lidc': h5py.File('/data_4t/Kaggle/backup/lidc/vol.hdf5', 'r'),
        'kaggle_mali': h5py.File('/data_4t/Kaggle/backup/kaggle/vol.hdf5', 'r'),
        'kaggle_beni': h5py.File('/data_4t/Kaggle/backup/kaggle/vol.hdf5', 'r'),
        'kaggle_testset': h5py.File('/data_4t/Kaggle/backup/kaggle/vol.hdf5', 'r'),
        'spie': h5py.File('/data_4t/Kaggle/backup/spie/vol.hdf5', 'r'),
    }

    for i in range(len(candis)):
        scan_id = candis[i]['scanid']
        source = candis[i]['source']

        # if cmp(scan_id,scan_id_tmp) !=0:
        if scan_id != scan_id_tmp:
            # if scan_id_tmp[0] == 'L':
            #     spac = v1[scan_id_tmp].attrs['spacing']
            #     lab = l1[scan_id_tmp]
            #     cen = []
            #     radius = []
            #     malign = []
            #     for ilab in range(len(lab)):
            #         cen.append(lab[ilab]['center'])
            #         radius.append(lab[ilab]['radius'])
            #         malign.append(lab[ilab]['attrs']['malignancy'])
            #
            # else:
            #     spac = v2[scan_id_tmp].attrs['spacing']
            #     lab = l2[scan_id_tmp]
            #     cen = []
            #     radius = []
            #     malign = []
            #     for ilab in range(len(lab)):
            #         cen.append(lab[ilab]['center'])
            #         radius.append(lab[ilab]['radius'])
            #         malign.append(lab[ilab]['malign'])

            spac = v[source_tmp][scan_id_tmp].attrs['spacing']

            keep, rf_center = nms(center_all, score_all, thre_dist, spac, 1)

            for ikeep in range(len(keep)):
                final_candis.append(candis[index_all[keep[ikeep]]])

            # evaluate performance
            # num_neg,num_pos,isFound=evaluate(rf_center,radius,cen,spac)
            # num_pos_t=num_pos_t+num_pos
            # num_neg_t=num_neg_t+num_neg
            # find_nod=find_nod+sum(isFound)
            # all_nod=all_nod+isFound.shape[0]
            # malign=np.array(malign)
            # # print scan_id_tmp
            # if scan_id_tmp[0]=='L':
            # 	find_ma=find_ma+sum(isFound*((malign>3)*1))
            # 	all_ma=all_ma+sum((malign>3)*1)
            # else:
            # 	find_ma=find_ma+sum(isFound*((malign>6)*1))
            # 	all_ma=all_ma+sum((malign>6)*1)
            # print iscan,scan_id_tmp,find_nod,all_nod,find_ma,all_ma,num_pos_t/(1.0+iscan),num_neg_t/(1.0+iscan)

            iscan = iscan + 1
            center_all = []
            score_all = []
            index_all = []
            scan_id_tmp = scan_id
            source_tmp = source

        if candis[i]['prob'] >= thre_score:
            center_all.append(candis[i]['center'])
            score_all.append(candis[i]['prob'])
            index_all.append(i)

        if i == len(candis) - 1:
            keep, rf_center = nms(center_all, score_all, thre_dist, spac, 1)
            for ikeep in range(len(keep)):
                final_candis.append(candis[index_all[keep[ikeep]]])

    return final_candis


def main(infile, outfile):
    candis = read_pkl_file(infile)
    candidate_list = main_nms(candis)
    print(len(candidate_list))
    write_pkl_file(outfile, candidate_list)

if __name__ == '__main__':
    # candis = np.load('/home/liax/py-faster-rcnn/kaggle/3dcnn_lidc_kaggle_candidates_v9_nomerge.npy')
    # candidate_list = main_nms(candis)
    # print(len(candidate_list))
    # np.save('/home/liax/py-faster-rcnn/kaggle/3dcnn_lidc_kaggle_candidates_v9_nomerge_nms.npy',candidate_list)
    #
    # main(
    #     infile='/data_4t/Kaggle/candidates/3dcnn_cddv9onv13_lidc_kaggle_candidates_v13_nomerge.pkl',
    #     outfile='/data_4t/Kaggle/candidates/3dcnn_cddv9onv13_lidc_kaggle_candidates_v13_nomerge_nms.pkl'
    # )
    # main(
    #     infile='/data_4t/Kaggle/candidates/3dcnn_cddv9onv13_kaggle_beni_candidates_v8_nomerge.pkl',
    #     outfile='/data_4t/Kaggle/candidates/3dcnn_cddv9onv13_kaggle_beni_candidates_v8_nomerge_nms.pkl'
    # )
    # main(
    #     infile='/data_4t/Kaggle/candidates/3dcnn_cddv9onv13_kaggle_testset_candidates_v8_nomerge.pkl',
    #     outfile='/data_4t/Kaggle/candidates/3dcnn_cddv9onv13_kaggle_testset_candidates_v8_nomerge_nms.pkl'
    # )
    # main(
    #     infile='/data_4t/Kaggle/candidates/3dcnn_cddv9onv13_spie_candidates_v8_nomerge.pkl',
    #     outfile='/data_4t/Kaggle/candidates/3dcnn_cddv9onv13_spie_candidates_v8_nomerge_nms.pkl'
    # )

    main(
        infile='/data_4t/Kaggle/candidates/3dcnn_cddv15_v2_kaggle_beni_candidates_v8_nomerge.pkl',
        outfile='/data_4t/Kaggle/candidates/3dcnn_cddv15_v2_kaggle_beni_candidates_v8_nomerge_nms.pkl'
    )
    main(
        infile='/data_4t/Kaggle/candidates/3dcnn_cddv15_v2_kaggle_testset_candidates_v8_nomerge.pkl',
        outfile='/data_4t/Kaggle/candidates/3dcnn_cddv15_v2_kaggle_testset_candidates_v8_nomerge_nms.pkl'
    )
    main(
        infile='/data_4t/Kaggle/candidates/3dcnn_cddv15_v2_spie_candidates_v8_nomerge.pkl',
        outfile='/data_4t/Kaggle/candidates/3dcnn_cddv15_v2_spie_candidates_v8_nomerge_nms.pkl'
    )
